#include "stdafx.h"
#include "SyscallMonForm.h"
#include "DriverLoader.h"
#include "ProcessMgr.h"
#include "EventMgr.h"
#include "util.h"

#include "../Shared/Protocol.h"

#include <string>
#include <shlwapi.h>
#pragma comment(lib,"shlwapi.lib")

CDriverLoader m_Driver;

extern HANDLE m_hQuitEvent;

bool ParseMessage(PUCHAR data);

BOOL LoadDriver(void)
{
	TCHAR szDirectory[MAX_PATH];
	CString DriverFilePath, DriverFilePathSymLink;
	LPCTSTR szDriverFileName = Is64Bit() ? _T("SyscallMon64.sys") : _T("SyscallMon32.sys");

	if (!AdjustPrivilege(SE_DEBUG_NAME))
	{
		MessageBox(NULL, L"Failed to get SE_DEBUG_NAME privilege!", L"Fatal Error", MB_OK | MB_ICONERROR);
		return FALSE;
	}

	if (!AdjustPrivilege(SE_LOAD_DRIVER_NAME))
	{
		MessageBox(NULL, L"Failed to get SE_LOAD_DRIVER_NAME privilege!", L"Fatal Error", MB_OK | MB_ICONERROR);
		return FALSE;
	}

	GetModuleFilePath(NULL, szDirectory, MAX_PATH);

	DriverFilePath.Format(L"%s\\%s", szDirectory, szDriverFileName);

	if (!PathFileExists(DriverFilePath))
	{
		CString err;
		err.Format(L"Could not found %s", szDriverFileName);

		MessageBox(NULL, err, L"Error", MB_OK | MB_ICONERROR);
		return FALSE;
	}

	DriverFilePathSymLink.Format(L"\\??\\%s", DriverFilePath);

	conn_context_data conn;
	conn.txsb = 'TXSB';
	conn.ver = 1;

	if (!m_Driver.Connect(L"\\SyscallMonPort", &conn, sizeof(conn)))
	{
		if (!m_Driver.Install(DriverFilePathSymLink, L"SyscallMon", L"SyscallMon"))
		{
			MessageBox(NULL, m_Driver.m_szErrorInfo, L"Fatal Error", MB_OK | MB_ICONERROR);
			return FALSE;
		}
		if (!m_Driver.Connect(L"\\SyscallMonPort", &conn, sizeof(conn)))
		{
			MessageBox(NULL, m_Driver.m_szErrorInfo, L"Fatal Error", MB_OK | MB_ICONERROR);
			return FALSE;
		}
	}
	else
	{
		lstrcpy(m_Driver.m_pServiceName, L"SyscallMon");
		lstrcpy(m_Driver.m_pDisplayName, L"SyscallMon");
	}

	return TRUE;
}

void UnloadDriver(void)
{
	m_Driver.Disconnect();
	m_Driver.Uninstall();
}

DWORD WINAPI DriverMessageRoutine(LPVOID lpContext)
{
	BOOL bIs64Bit = Is64Bit();
	UCHAR buf[4096];

	HANDLE hEvents[2];

	hEvents[0] = m_ProcessMgr.m_hReadyEvent;
	hEvents[1] = m_EventMgr.m_hReadyEvent;
	WaitForMultipleObjects(2, hEvents, TRUE, INFINITE);

	OVERLAPPED ovlp;
	ovlp.hEvent = ::CreateEvent(NULL, FALSE, FALSE, NULL);
	ovlp.Internal = 0;
	ovlp.InternalHigh = 0;
	ovlp.Pointer = 0;
	ovlp.Offset = 0;
	ovlp.OffsetHigh = 0;

	hEvents[0] = ovlp.hEvent;
	hEvents[1] = theApp.m_hQuitEvent;

	while (1)
	{
		if (WaitForSingleObject(theApp.m_hQuitEvent, 0) == WAIT_OBJECT_0)
			break;

		memset(buf, 0, sizeof(buf));
		HRESULT hres = m_Driver.Read(buf, sizeof(buf), &ovlp);
		if (hres == HRESULT_FROM_WIN32(ERROR_IO_PENDING) || hres == S_OK)
		{
			if (hres == HRESULT_FROM_WIN32(ERROR_IO_PENDING))
			{
				if (WaitForMultipleObjects(2, hEvents, FALSE, INFINITE) == (WAIT_OBJECT_0 + 1))//time to quit
					break;
			}

			PUCHAR dataBuf = buf + (bIs64Bit ? 16 : 12);

			ParseMessage(dataBuf);
		}
		else//if driver errors, wait 1 sec and try again
		{
			Sleep(1000);
		}
	}

	CloseHandle(ovlp.hEvent);

	return 0;
}

DWORD AddIgnoreProcess(DWORD dwProcessId)
{
	cls_pid_data data;
	DWORD dwParentProcessId = 0;
	DWORD dwReadBytes = 0;

	data.protocol = cls_add_ignore_pid;
	data.ProcessId = dwProcessId;

	return m_Driver.Send(&data, sizeof(data), NULL, 0, &dwReadBytes);
}

ULONG64 GetKeSystemTime(void)
{
	ULONG64 time = 0;
	DWORD dwReadBytes = 0;

	UCHAR protocol[1];
	protocol[0] = cls_get_system_time;

	if (m_Driver.Send(protocol, sizeof(protocol), &time, sizeof(time), &dwReadBytes))
		return time;

	return 0;
}

BOOL GetProcessImagePath(DWORD dwProcessId, LPWSTR szImagePath, DWORD dwImagePathLength)
{
	cls_pid_data data;
	DWORD dwReadBytes = 0;

	data.protocol = cls_get_process_path;
	data.ProcessId = dwProcessId;

	memset((PVOID)szImagePath, 0, dwImagePathLength);

	return m_Driver.Send(&data, sizeof(data), (LPVOID)szImagePath, dwImagePathLength - sizeof(WCHAR), &dwReadBytes);
}

BOOL GetImagePathByAddress(DWORD dwProcessId, ULONG64 ImageBase, LPWSTR szImagePath, DWORD dwImagePathLength)
{
	cls_get_image_path_data data;
	DWORD dwReadBytes = 0;

	data.protocol = cls_get_image_path;
	data.ProcessId = dwProcessId;
	data.ImageBase = ImageBase;

	memset((PVOID)szImagePath, 0, dwImagePathLength);

	return m_Driver.Send(&data, sizeof(data), (LPVOID)szImagePath, dwImagePathLength - sizeof(WCHAR), &dwReadBytes);
}

DWORD GetParentProcessId(DWORD dwProcessId)
{
	cls_pid_data data;
	DWORD dwParentProcessId = 0;
	DWORD dwReadBytes = 0;

	data.protocol = cls_get_parent_pid;
	data.ProcessId = dwProcessId;

	if (m_Driver.Send(&data, sizeof(data), &dwParentProcessId, sizeof(dwParentProcessId), &dwReadBytes))
		return dwParentProcessId;

	return 0;
}

BOOL GetProcessCreateTime(DWORD dwProcessId, FILETIME *CreateTime)
{
	cls_pid_data data;
	DWORD dwReadBytes = 0;

	data.protocol = cls_get_process_createtime;
	data.ProcessId = dwProcessId;

	return m_Driver.Send(&data, sizeof(data), CreateTime, sizeof(FILETIME), &dwReadBytes);
}

BOOL GetProcessBaseInfo(DWORD dwProcessId, DWORD *ParentProcessId, FILETIME *CreateTime, BOOLEAN *bIs64Bit, DWORD *SessionId, std::wstring &ImagePath)
{
	DWORD dwReadBytes = 0;

	cls_pid_data data;

	data.protocol = cls_get_process_baseinfo;;
	data.ProcessId = dwProcessId;

	cls_get_process_baseinfo_data out = { 0 };

	if (m_Driver.Send(&data, sizeof(data), &out, sizeof(out), &dwReadBytes))
	{
		*ParentProcessId = out.ParentProcessId;
		*CreateTime = *(FILETIME *)&out.CreateTime;
		*bIs64Bit = out.Is64Bit;
		*SessionId = out.SessionId;
		ImagePath = out.ImagePath;		
		return TRUE;
	}
	return FALSE;
}

void SetEnableCapture(BOOL bEnable)
{
	cls_set_capture_data data;

	data.protocol = cls_set_capture;
	data.Capture = bEnable;

	BOOL b = m_Driver.Send(&data, sizeof(data), NULL, 0, NULL);
}